Amazon SageMaker上でText Generation InferenceのMessages APIを使って、オープンなLLMをOpenAI互換なAPIで使ってみた。

nokomoro3

2024.02.24

こんちには。

データアナリティクス事業本部機械学習チームの中村(nokomoro3)です。

今回は、Hugging Faceが作っているText Generation InferenceのMessages API機能を、Amazon SageMaker上で使ってEndpointを作成することで、OpenAIのLLMのようなAPIで呼び出してみたいと思います。

Text Generation Inferenceとは

Text Generation Inferenceは、LLMをデプロイして利用するためのツールキットです。

以下のような多くの最適化と機能を実装しています。

プロダクション対応（Open Telemetryによる分散トレース、Prometheusメトリクス）
複数のGPUで推論を高速化するテンソル並列処理
bitsandbytesとGPT-Qによる量子化
Stop Sequences
Fine-tuningのサポート

Text Generation Inferenceは、1.4.0からMessages APIというものが実装されています。

From OpenAI to Open LLMs with Messages API on Hugging Face

Messages APIはOpenAIのChatCompletion APIと互換性のあるAPIを提供します。

Messages APIにより、顧客や開発者はOpenAIモデルからオープンLLMへシームレスに移行することが可能となっており、このAPIはOpenAIのクライアントライブラリや、LangChainやLlamaIndexのようなサードパーティツールで直接使用することができます。

今回はこちらがAmazon SageMaker上でのデプロイにも対応していることがわかったので、そちらを試してみたいと思います。

Messages API

ドメインの作成

Amazon SageMakerのコンソールから、ドメインをクリックします。

ドメイン一覧が表示されますので、「ドメインを作成」を押下します。

ドメイン作成方法の選択

「シングルユーザー向けの設定(クイックセットアップ)」を選択して、「設定」を押下します。

参考：作成されるIAMロールとS3バケット

作成されるIAMロールには以下のマネージドポリシーがアタッチされています。

AmazonSageMakerCanvasAIServicesAccess
AmazonSageMakerCanvasFullAccess
AmazonSageMakerFullAccess

また、以下のようなカスタムポリシーもアタッチされています。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject",
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::*"
            ]
        }
    ]
}

またSageMakerが使用するS3バケットが以下の命名規則で必要となっており、存在しない場合は作成される形となります。

s3://sagemaker-{リージョン名}-{アカウントID}

作業環境の立ち上げ

ドメイン上に作業環境として、JupyterLabを立ち上げます。

Studioの立ち上げ

ドメイン作成後は、以下のような画面に遷移するので、起動メニューから「Studio」を押下します。

JupyterLabのスペース作成

Studio起動後は以下のような画面となるので、JupyterLabをクリックします。

「Create JupyterLab Space」を押下します。

名前を入力して、「Create space」を押下します。

今回はjupyuterlab-2024-02-24という名前としました。

JupyterLabの立ち上げ

以下のような画面に遷移するので、「Run space」を押下します。

この「Run space」時にJupyterLabを動かすインスタンスと、使用するコンテナイメージを指定することができます。

今回はインスタンスをml.t3.medium、コンテナイメージをSageMaker Distribution 1.4のデフォルトで実行しています。

起動状態となったら、「Open JupyterLab」を押下します。

JupyterLabが起動しますので、以下をクリックしてPythonのNotebookを起動します。

以降はこのNotebook上で作業をします。

Notebook上での作業

Notebook上でのコードは以下を参考とします。

Messages API

JupyterLab上での準備

まずはPythonのバージョンを確かめておきます。

!python --version

# Python 3.10.13

次に、コンテナ内のSageMaker SDK関連のパッケージをアップデートします。

!pip install -U sagemaker

# (途中略)
# Successfully installed boto3-1.34.49 botocore-1.34.49 dill-0.3.8 docker-7.0.0 multiprocess-0.70.16 s3transfer-0.10.0 sagemaker-2.209.0

各種インポートを実行します。

import json
import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

SageMaker実行ロールを取得します。

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

print(f"{role=}")

# sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /home/sagemaker-user/.config/sagemaker/config.yaml
# role='arn:aws:iam::{AWSアカウント番号}:role/service-role/SageMaker-ExecutionRole-20240224T154483'

Endpointの作成

Hugging Face Hubで公開されているモデルIDやその設定を実行します。

今回はモデルIDとしてmistralai/Mistral-7B-Instruct-v0.2を使ってみます。

mistralai/Mistral-7B-Instruct-v0.2 · Hugging Face

# Hub Model configuration. https://huggingface.co/models
hub = {
    'HF_MODEL_ID':'mistralai/Mistral-7B-Instruct-v0.2',
    'SM_NUM_GPUS': json.dumps(1),
    'MESSAGES_API_ENABLED': json.dumps(True), # ここにjson.dumpsを追加
}

MESSAGES_API_ENABLEDの設定はサンプル通りではうまく行かなかったためjson.dumpsでTrueを囲っています。

次にモデルをインスタンス化します。バージョン1.4.0はText Generation Inferenceのバージョンとなります。

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=get_huggingface_llm_image_uri("huggingface", version="1.4.0"),
    env=hub,
    role=role,
)

# sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /home/sagemaker-user/.config/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /home/sagemaker-user/.config/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml
# sagemaker.config INFO - Not applying SDK defaults from location: /home/sagemaker-user/.config/sagemaker/config.yaml

最後にEndpointのデプロイを実行します。今回はEndpointのインスタンスとしてml.g4.2xlargeを選択します。

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    container_startup_health_check_timeout=300,
)

Endpointでの推論テスト

英語でチャットを試してみます。

# send request
response = predictor.predict({
    "messages": [
        {"role": "user", "content": "What is deep learning?"}
    ],
    "model": "mistralai/Mistral-7B-Instruct-v0.2",
    "max_tokens": 1024
})

"role": "systems"を入れるとこのモデルの場合はテンプレートに合わないようでエラーとなりました。

"model"は参考にしたコードには含まれていませんでしたが、指定しないとエラーとなったので指定しています。

また"max_tokens"を入れないと、レスポンスが切れてしまったため指定をしています。

結果は以下の通りで、レスポンスが返ってくることが確認できました。

print(response["choices"][0]["message"]["content"])

 Deep learning is a subfield of machine learning, which in turn is a subset of artificial intelligence (AI). 
 Deep learning models are designed to recognize patterns and learn representations of data at multiple levels of abstraction, 
 mimicking the way the human brain functions. These models are able to do so by using deep neural networks, 
 which consist of interconnected layers of artificial neurons, each taking in data from the previous layer and 
 passing on a transformed representation. Deep learning models can learn unsupervised from unstructured or unlabeled data, 
 as well as supervised from labeled data, and have achieved impressive results in various applications such as image and 
 speech recognition, natural language processing, autonomous driving, and many others. They have been particularly 
 successful due to the availability of large amounts of data and computational power, making it possible to train complex 
 deep neural networks that can learn and represent complex patterns in data.

Endpointの確認

Endpointは、Studioの画面から「Deployments」->「Endpoints」を選択すると確認することができます。

こちらでEndpointsの立ち上げには、6分程度かかったことが分かります。

CloudWatchのロググループは、このEndpoint名を使って/aws/sagemaker/Endpoints/{Endpoint名}となっています。

エラーなどが発生し、詳細を確認したい場合はこのロググループを参照ください。Endpoint立ち上げ時のエラーも確認することができます。

リソースの削除

最後に費用が掛からないようリソースを削除しておきます。

先ほどのEndpoint画面から選択し、「Delete」を押下すると削除することが可能です。

またJupyterLabも課金の対象となりますので、以下から停止しておきましょう。

最後に必要であればドメインのユーザを削除し、ドメインを削除します。それぞれユーザの編集とドメインの編集から削除もできます。

また、IAMロールも残したくなければそちらも削除されてください。

トラブルシューティング

モデルがダウンロードできない

今回は必要ありませんが、指定するモデルIDによっては利用規約への同意が必要なLLMがあり、Hugging Faceのアカウントを作成して、トークンを生成しておく必要があります。

以下のような画面が出るモデルがその対象となります。（例はgoogle/gemma-7b-it）

ログインまたはサインアップ後、アカウント管理画面でAccess Tokenを作成します。

https://huggingface.co/settings/tokens

トークンを作成後は、以下のHugging Face Hubの設定にHUGGING_FACE_HUB_TOKENを加えます。

# Hub Model configuration. https://huggingface.co/models
hub = {
    'HF_MODEL_ID':'pfnet/plamo-13b-instruct',
    'SM_NUM_GPUS': json.dumps(1),
    'MESSAGES_API_ENABLED': json.dumps(True),
    'HUGGING_FACE_HUB_TOKEN': '{Hugging Face Hubのトークン}'
}

サポート外のモデルID

サポート外のモデルIDを'HF_MODEL_ID'に指定すると、deploy時に以下のエラーが発生します。

（以下はCloudWatchに出力されたエラーで確認）

> File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 196, in serve_inner
    model = get_model(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/__init__.py", line 422, in get_model
    raise ValueError(f"Unsupported model type {model_type}")
ValueError: Unsupported model type gemma

サポートしているモデル一覧は以下などを参照ください。

なお、Text Generation Inferenceの1.4.2ではgoogle/gemma-7b-itに対応しており、以下のようにすれば解決するかと思い試したのですが、

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=get_huggingface_llm_image_uri("huggingface", version="1.4.2"), # 1.4.2に変更
    env=hub,
    role=role,
)

sagemaker-2.209.0がまだtext-generation-inferenceの1.4.2に対応していないため使用ができない旨のメッセージとなりました。

SageMaker SDKのアップデートを今後待ちたいと思います。

chat_templateが無いエラー

サポート対象となっている'HF_MODEL_ID'だったとしても、chat_templateが無いモデルに対して'MESSAGES_API_ENABLED'を有効としてテキスト生成をさせようとするとエラーとなるようです。

ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received client error (422) from primary with message
"{"error":"Template error: template not found","error_type":"template_error"}".

上記はmiscrosoft/phi-2などのモデルで発生しました。確かにtokenizer_config.jsonにchat_templateが無いようです。

tokenizer_config.json · microsoft/phi-2 at main

単にchat_templateが書いてないだけのケースもありますので、Hubの更新を待つかそれ以外のチャットに適したモデルを選択するようにしましょう。

chat_templateにマッチしないエラー

こちらは今回のmistralai/Mistral-7B-Instruct-v0.2で発生しました。

ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received client error (422) from primary with message
"{"error":"Template error: syntax error: Conversation roles must alternate user/assistant/user/assistant/... (in <string>:1)","error_type":"template_error"}"

以下のようにmessagesに"role": "system"を加えると、このモデルの場合エラーとなるようです。

# send request
response = predictor.predict({
    "messages": [
        {"role": "system", "content": "You are a helpful assistant." },
        {"role": "user", "content": "What is deep learning?"}
    ],
    "model": "mistralai/Mistral-7B-Instruct-v0.2",
    "max_tokens": 1024
})

このようなエラーが出た場合は、よく読んで個別に対処をする必要がありそうです。